import cv2
import math
import numpy as np
from math_utility import get_distance


# 计算核函数矩阵
def get_kernel(ks):
    """
    Args:
        ks: 核函数的尺寸|size of kernel

    Returns: kernel模板|kernel
    """
    kernel = np.zeros([ks, ks])
    center = int((ks - 1) / 2)   # 中心点坐标
    for i in range(0, center + 1):
        for j in range(0, i + 1):
            x = get_distance((i, j), (center, center))  # x为矩阵中两个元素的距离
            y = 2 - math.pow(2 * x / (ks - 1), 4)    # y为矩阵元素的实际值，由幂函数定义
            if y >= 1:
                y = round(y, 2)
                # 对称法直接赋值
                kernel[i][j] = y
                kernel[j][i] = y
                kernel[i][ks - 1 - j] = y
                kernel[ks - 1 - i][j] = y
                kernel[j][ks - 1 - i] = y
                kernel[ks - 1 - j][i] = y
                kernel[ks - 1 - i][ks - 1 - j] = y
                kernel[ks - 1 - j][ks - 1 - i] = y
    return kernel


# 正规化至0-255
def normalize(matrix):
    """
    Args:
        matrix: 二维矩阵

    Returns: 灰度图像

    """
    max_value = np.max(matrix)
    min_value = np.min(matrix)
    matrix = (matrix - min_value) / max_value * 255
    return matrix


# 获取被聚类后的点集
def get_maxlist(matrix, ks, max_length):
    """
    Args:
        matrix: 二维矩阵
        ks: kernel尺寸
        max_length: 取得前max_length个最大值

    Returns: max_length个最大值所对应的中心点

    """
    center = []
    for i in range(max_length):
        value = np.argmax(matrix)
        x = int(value / matrix.shape[1])
        y = value % matrix.shape[1]
        # 已经不存在任何值
        if x == 0 and y == 0:
            break
        center.append((x, y))
        x1, x2, y1, y2 = x - ks, x + ks, y - ks, y + ks
        if x - ks < 0:
            x1 = 0
        if y - ks < 0:
            y1 = 0
        if x + ks >= matrix.shape[0]:
            x2 = matrix.shape[0] - 1
        if y + ks >= matrix.shape[1]:
            y2 = matrix.shape[1] - 1
        matrix[x1: x2, y1: y2] = 0
    return center


# 幂函数加权聚类法
def weighted_cluster(matrix, point_set, ks=15, max_length=20):
    """
    Args:
        matrix: 二维矩阵
        point_set: 被聚类的点集
        ks: 核函数的尺寸，为奇数时聚类效果最好
        max_length: 输出的最大聚类数量

    Returns: 热力图和聚类后的点集
    """
    if ks > matrix.shape[0] or ks > matrix.shape[1]:
        raise ValueError('The ks can not bigger than the width of matrix.')
    elif ks <= 0:
        raise ValueError('The ks can not less than or equal to zero.')
    kernel = get_kernel(ks)     # 获取kernel矩阵
    radius = int((ks - 1) / 2)  # 计算kernel半径
    for i in range(len(point_set)):
        center = point_set[i]
        x1 = center[0] - radius
        x2 = center[0] + radius + 1
        y1 = center[1] - radius
        y2 = center[1] + radius + 1
        # 判断边界情况，计算偏移量
        dx1, dy1, dx2, dy2 = 0, 0, 0, 0
        if x1 < 0:
            dx1 = radius - center[0]
            x1 = 0
        if y1 < 0:
            dy1 = radius - center[1]
            y1 = 0
        if x2 >= matrix.shape[0]:
            dx2 = radius + center[0] - matrix.shape[0] + 2
            x2 = matrix.shape[0] - 1
        if y2 >= matrix.shape[1]:
            dy2 = radius + center[1] - matrix.shape[1] + 2
            y2 = matrix.shape[1] - 1
        matrix[x1: x2, y1: y2] += kernel[dx1: kernel.shape[0] - dx2, dy1: kernel.shape[1] - dy2]    # 将kernel值累加到热图上
    matrix = normalize(matrix)  # 归一化至0-255
    matrix = np.array(matrix, dtype=np.uint8)
    matrix_copy = matrix.copy()
    point_set = get_maxlist(matrix_copy, ks, max_length)
    return matrix, point_set


if __name__ == '__main__':
    frame = np.zeros([400, 480])
    ps = [[35, 23], [35, 24], [34, 19], [34, 17], [36, 19], [36, 22], [36, 23], [36, 25], [37, 17],
          [44, 155], [45, 153], [44, 160], [48, 153], [46, 140], [88, 77], [88, 78], [85, 85], [84, 76],
          [237, 198], [235, 199], [232, 190], [239, 188], [235, 201], [237, 202], [199, 198], [233, 191],
          [312, 43], [311, 42], [312, 48], [311, 47], [310, 41], [309, 42], [313, 53], [216, 77], [251, 280],
          [57, 178], [95, 95], [177, 148], [175, 144], [173, 150], [284, 188], [251, 78], [153, 140], [58, 155],
          [355, 409], [357, 408], [357, 405], [352, 406], [354, 411], [354, 402], [358, 400], [358, 413], [360, 410],
          [225, 356], [224, 356], [220, 356], [220, 358], [221, 351], [229, 349], [215, 355], [219, 353], [222, 350]]
    heat_map, centers = weighted_cluster(frame, ps, 17, 20)
    for i in range(len(centers)):
        print(i)
        print('center:', centers[i])
    cv2.imshow('heat map', heat_map)
    cv2.waitKey(0)
